/*
  Copyright (c) 2008 Florent D'halluin , Sylvain Calinon, 
  LASA Lab, EPFL, CH-1015 Lausanne, Switzerland, 
  http://www.calinon.ch, http://lasa.epfl.ch

  The program is free for non-commercial academic use. 
*/

#include "libxml/parser.h"
#include "libxml/tree.h"
#include <cstdlib>
#include <cstring>
#include "IK.h"
#include <math.h>
#include <iostream>
#include <fstream>

Matrix3 generate_rotation(short axe,float angle)
{
    Matrix3 r;
    if(axe== 0)
      {
	r.RotationX(angle);
      }
    else if(axe== 1)
      {
	r.RotationY(angle);
      }
    else
      {
	r.RotationZ(angle);
      }
    return r;
}

Matrix3 diff_rotation(short axe,float angle)
{
  Matrix3 r;
    if(axe== 0)
      {
	r.RotationX(angle + M_PI/2);
	r(0,0) = 0.;
      }
    else if(axe== 1)
      {
	r.RotationY(angle + M_PI/2);
	r(1,1) = 0.;
      }
    else
      {
	r.RotationZ(angle + M_PI/2);
	r(2,2) = 0.;
      }
  
  return r;
}

Matrix4 Segment::get_matrix(float angle)
{
  Matrix4 hom = Matrix4::IDENTITY;
  Matrix3 r = generate_rotation(axe,angle);
  hom.SetOrientation(r);
  hom.SetTranslation(seg);
  return hom;
}

Matrix4 Segment::get_diff_mat(float angle)
{
  Matrix4 hom; 
  Matrix3 r = diff_rotation(axe,angle);
  hom.SetOrientation(r);
  return hom;
}

Quaternion Segment::get_quaternion(float angle)
{
  Quaternion q;
  if(axe == 0)
    {
      Quaternion::RotationX(angle,q);
    }
  if(axe == 1)
    {
      Quaternion::RotationY(angle,q);
    }
  if(axe == 2)
    {
      Quaternion::RotationZ(angle,q);
    }
  return q;
}

Quaternion Segment::get_quaternion_speed(float angle)
{
  Quaternion q = Quaternion::ZERO;
  q.w = -sin(angle * .5f)*.5f;
  if(axe == 0)
    q.x = cos(angle *.5f)*.5f;
  if(axe == 1)
    q.y = cos(angle *.5f)*.5f;
  if(axe == 2)
    q.z = cos(angle *.5f)*.5f;
  return q;
}

float Segment::getJointLimitGradient(float angle)
{
  float g;
  if(this->max_joint_angle == this->min_joint_angle)
    {
      std::cout << "joints limits not defined .. " << std::endl;
      return 0; // NaN
    }
  else
    {
      g = (max_joint_angle-min_joint_angle)* (max_joint_angle-min_joint_angle);
      g *= (2*angle-max_joint_angle-min_joint_angle);
      g /= 4*(max_joint_angle-angle)*(max_joint_angle-angle);
      g/= (angle-min_joint_angle)*(angle-min_joint_angle);
      return g;
    }
}

IKModel::IKModel(int * axes,Vector3 * segs_v, int seg_num, Vector3 end_vec)
{
  segs = new Segment*[seg_num];
  for(int i=0;i<seg_num;i++)
    {
      std::cout << axes[i] << std::endl;
      segs_v[i].Print();
      segs[i] = new Segment(axes[i],segs_v[i]);
    }
  end_vec.Print();
  end_vector = end_vec;
  this->seg_num = seg_num;
}


IKModel::IKModel(const char filename[])
{
  xmlDoc * doc;
  xmlNode * root;
  xmlNode * cur_node = NULL;
  xmlNode * cur_seg_e = NULL;
  int n_seg =0;
  doc = xmlReadFile(filename, NULL, 0);
  char * axe_xml;
  char * x_c, * y_c, * z_c;
  int axe_i[20];
  float max[20],min[20];
  int limits_flag[20];
  for(int i=0;i<20;i++) limits_flag[i]=0;
  Vector3 seg_xml[20];
  Vector3 end_effector;
  root = xmlDocGetRootElement(doc);
  for (cur_node = root->children; cur_node; cur_node = cur_node->next) {
    if (cur_node->type == XML_ELEMENT_NODE) {
      if( strcmp((char*)cur_node->name,"segment")==0) 
	{
	  axe_xml = (char *) xmlGetProp(cur_node,(xmlChar *)"axe");
	  x_c = (char *) xmlGetProp(cur_node,(xmlChar *)"x");
	  y_c = (char *) xmlGetProp(cur_node,(xmlChar *)"y");
	  z_c = (char *) xmlGetProp(cur_node,(xmlChar *)"z");
	  axe_i[n_seg] = atoi(axe_xml);
	  seg_xml[n_seg](0) = atof(x_c);
	  seg_xml[n_seg](1) = atof(y_c);
	  seg_xml[n_seg](2) = atof(z_c);
	  
	  for(cur_seg_e = cur_node->children;cur_seg_e;
	      cur_seg_e = cur_seg_e->next)
	    {
	      if( strcmp((char*)cur_seg_e->name,"joint_limits")==0)  
		{
		  max[n_seg] = atof((char*)
				    xmlGetProp(cur_seg_e,(xmlChar *)"max"));
		  min[n_seg] = atof((char*)
				    xmlGetProp(cur_seg_e,(xmlChar *)"min"));
		  // std::cout << max[n_seg] << std::endl;
		  limits_flag[n_seg] = 1;
		}
	    }
	  n_seg++;
	}
      if( strcmp((char*)cur_node->name,"end")==0)
	{
	  x_c = (char *) xmlGetProp(cur_node,(xmlChar *)"x");
	  y_c = (char *) xmlGetProp(cur_node,(xmlChar *)"y");
	  z_c = (char *) xmlGetProp(cur_node,(xmlChar *)"z");
	  end_effector(0) = atof(x_c);
	  end_effector(1) = atof(y_c);
	  end_effector(2) = atof(z_c);
	}
    }
    
  }
  xmlFreeDoc(doc);

  this->segs = new Segment*[n_seg];
  for(int i=0;i<n_seg;i++)
    {
      this->segs[i] = new Segment(axe_i[i],seg_xml[i]);
      if(limits_flag[i])
	{
	  this->segs[i]->max_joint_angle = max[i];
	  this->segs[i]->min_joint_angle = min[i];
	}
    }
  this->end_vector = end_effector;
  this->seg_num = n_seg;
  this->f = new std::ofstream("joints_trajectory.txt");
  this->joint_positions = new Vector[seg_num+1];

}

void IKModel::dump(void)
{
  std::cout << "Model Dump" << std::endl;
  for(int i=0;i<seg_num;i++)
    {
      std::cout << "*************" << std::endl;
      std::cout << segs[i]->axe << std::endl;
      segs[i]->seg.Print();
    }
  end_vector.Print();
}

Vector3 IKModel::direct(Vector angles)
{
  Matrix4 p = Matrix4::IDENTITY;
  Matrix4 h;
  for(int i=0;i<seg_num;i++)
    {
      h =segs[i]->get_matrix(angles[i]); 
      //h.Print();
      p = p*h;
      joint_positions[i] = p.GetTranslation();
    }
  Vector3 result = p.Transform(end_vector);
  joint_positions[seg_num] = result;
  return result;
}

Vector IKModel::HMatrixToPosition(Matrix4 p)
{
  // convert an Homogenous Matrix transformation to (x,y,z,Phi1,Ph2,Ph3) 
  // of the end effector 
  Vector orientation,position;
  Vector result;
  Matrix3 R = p.GetOrientation();
  Quaternion q;
  Quaternion::FromMatrix3(R,q);
  // std::cout << q.Norm() << "--\n";

  if ( (q.w < 0) ) 
    {
      q = -q;
      }
  
  //float nfactor = (2./3.) * acos(q(0))/ q.GetSubVector(1,3).Norm();
  
  // float nfactor = 1.;
  
  orientation(0) = q.x;
  orientation(1) = q.y;
  orientation(2) = q.z;
  position = p.Transform(end_vector);
  result = position.GetSubVector(0,3);
  result.Resize(6);
  result.SetSubVector(3,orientation);
  return result;
}

Vector IKModel::direct_orientation(Vector angles)
{
  Matrix p(4,4);
  Matrix h;
  Matrix R;
  Vector position,orientation;
  Quaternion q = Quaternion::IDENTITY;
  Vector result(6);
  for(int i=0;i<4;i++)
    p(i,i)=1;
  for(int i=0;i<seg_num;i++)
    {
      h =segs[i]->get_matrix(angles[i]); 

      q *= segs[i]->get_quaternion(angles[i]);

      //h.Print();
      p = p*h;
      joint_positions[i] = p.GetColumn(3).GetSubVector(0,3);
    }
  position = p*end_vector;
  position = position.GetSubVector(0,3);
  position.Resize(6);
  position(3) = q.x;
  position(4) = q.y;
  position(5) = q.z;
  result = position;
  //result = HMatrixToPosition(p); 
  joint_positions[seg_num] = result.GetSubVector(0,3);
  return result;
}  

Matrix4 IKModel::getMatrix4(Vector angles)
{
  Matrix4 p;
  Matrix4 h;
  Matrix4 final;
  Vector q, position,orientation;
  for(int i=0;i<4;i++)
    {
      p(i,i)=1;
      final(i,i) = 1;
    }
  
  for(int i=0;i<seg_num;i++)
    {
      h =segs[i]->get_matrix(angles[i]); 
      //h.Print();
      p = p*h;
      //joint_positions[i] = p.GetColumn(3).GetSubVector(0,3);
    }
  final.SetTranslation(end_vector);
  p = p*final;
  return p;
}

Matrix IKModel::jacobian(Vector angles)
{
  /* compute the jacobian of the system, at the position given
     by anges values .. */
  Vector3 v;
  Matrix4 p;
  Matrix4 h;
  Matrix jacobian(seg_num,3);
  for(int i=0;i<seg_num;i++)
    {
      p = Matrix4::IDENTITY;
      for(int j=0;j<seg_num;j++)
	{
	  if(j==i)
	    {
	      h = segs[i]->get_diff_mat(angles[i]);
	    }
	  else
	    {
	      h = segs[j]->get_matrix(angles[j]);
	    }
	  p = p*h;
	}
      v = p.Transform(end_vector);
      jacobian.SetRow(v,i);
    }
  return jacobian.Transpose(); // hem ... 
}

Matrix IKModel::jacobian_orientation(Vector angles)
{
  Vector position;
  Vector orientation(3);
  Matrix p(4,4);
  Matrix r(3,3);
  Vector q;
  Matrix h;
  Matrix jacobian(seg_num,6);
  for(int i=0;i<seg_num;i++)
    {
      p.Zero();
      for(int k=0;k<4;k++) p(k,k)=1;
      for(int j=0;j<seg_num;j++)
	{
	  if(j==i)
	    {
	      h = segs[i]->get_diff_mat(angles[i]);
	    }
	  else
	    {
	      h = segs[j]->get_matrix(angles[j]);
	    }
	  p = p*h;
	}
      position = p*end_vector;
      position = position.GetSubVector(0,3);
      Quaternion q = Quaternion::IDENTITY;
      for(int j=0;j<seg_num;j++)
	{
	  if(j==i)
	    {
	      q *= segs[i]->get_quaternion_speed(angles[i]);
	    }
	  else
	    {
	      q *= segs[j]->get_quaternion(angles[j]);
	    }
	}       
      position.Resize(6);
      position(3) = q.x;
      position(4) = q.y;
      position(5) = q.z;
      jacobian.SetRow(position,i);
    }
  return jacobian.Transpose(); // hem ... 
}

int IKModel::inverse_orientation(Vector start,Vector otarget,Vector * next_angles)
{
  Matrix J,Jinv,Jtmp;
  Matrix Identity;
 
  Vector x,angles=start;
  x = direct_orientation(start);
  Vector dx,da;
  Vector ctarget;
  Vector current_x;
  dx = otarget - x;
  float alpha;
  float norm = dx.Norm() ;

  int n_iter = (int)floor(norm) + 1;

  for(int i = 1;i<= n_iter;i++)
    {
      alpha = (float) i/ (float) n_iter;
      ctarget = otarget*alpha + x*(1-alpha);
      angles = inverse_step_o(angles,ctarget,1,1);
    }
    
  x = direct_orientation(angles);
  dx = otarget-x;
  norm = dx.GetSubVector(3,3).Norm();
    std::cout << n_iter << " :: " << norm << std::endl;
    *next_angles = angles;
    return 1;
}


Vector IKModel::inverse_step_o(Vector start,Vector target,
			       float linSpeedThresh,float angSpeedThresh)
{
  Matrix J,Jinv,Jtmp;
  Matrix Identity;
  Matrix Weight(seg_num,seg_num);
  Matrix Weight_inv(seg_num,seg_num);
  Vector x,angles=start;
  Vector angles_t;
  Vector dx,dp,dor;
  Vector da;
  float prev_g[seg_num];
  Identity.Resize(this->seg_num,this->seg_num);
  for(int i=0;i<this->seg_num;i++) 
    {
      Identity(i,i) = 1;
      prev_g[i] = segs[i]->getJointLimitGradient(start(i));
    }
 
  
  Weight.Zero();
  Weight_inv.Zero();
  x = direct_orientation(angles);
  
  dx = (target-x);
  
  J = this->jacobian_orientation(angles);
  
  Jtmp = J*J.Transpose() + Identity*0.001;
  Jinv = J.Transpose()*Jtmp.Inverse();
  da = Jinv*dx;
  angles_t = angles+da;

  for(int i=0;i<seg_num;i++)
    {
      float g = segs[i]->getJointLimitGradient(angles_t(i));
      if((fabs(g)-prev_g[i]) > 0)
	{
	  Weight(i,i) = 1 + fabs(g); 
	  Weight_inv(i,i) = 1/(1+fabs(g));
	}
      else 
	{
	  Weight(i,i) = 1; 
	  Weight_inv(i,i) = 1;
	}
      prev_g[i] = fabs(g);
    }
  // Weight.Print();
  // Weighted Pseudo Inverse 
  Jtmp = J*Weight_inv*J.Transpose() + Identity*0.01; // + damping .. lambda*Identity .. 
  Jinv = Weight_inv *J.Transpose()*Jtmp.Inverse(); 
  //Jinv.Print();
  
  /*  
  Jtmp = J*J.Transpose();
  Jinv = J.Transpose() * Jtmp.Inverse();
  */  
  if(!Jtmp.IsInverseOk())
	{
	  std::cout << "can't inverse JtT" << std::endl;
	}
  
  da = Jinv*dx;
  angles = angles+da;
  return angles;

}



Vector IKModel::inverse_step(Vector currentJoints,Vector3 target,
			     float linSpeedThresh,float angSpeedTresh,
			     Vector refPosture)
{
  Matrix J,Jtmp,Jinv,Identity;
  Vector angles = currentJoints;
  int njoints = angles.Size();
  float norm;
  Vector3 x,dx;
  Vector da;
  Identity.Resize(this->seg_num,this->seg_num);
  for(int i=0;i<this->seg_num;i++) Identity(i,i)=1;

  x = this->direct(currentJoints);
  dx = (target-x);
  norm = dx.Norm();
  if(norm > linSpeedThresh)
    dx = dx * (linSpeedThresh/norm);

  J= this->jacobian(angles);
  
  // pseudo inverse .. 
  Jtmp = J* J.Transpose();
  Jinv = Jtmp.Inverse();
  
  if(!Jtmp.IsInverseOk())
    {
      std::cout << "cant inverse JtT" << std::endl;
      // here we should break .. 
    }
  Jinv = J.Transpose()*Jinv;
  
  da = Jinv*dx + (Identity-Jinv*J)*(refPosture - angles)*0.05;

  /*
  if(da.Norm() > angSpeedTresh)
    da = da*(angSpeedTresh/da.Norm()); */
  

  for(int i=0; i < this->seg_num;i++)
    {
      if(fabs(da(i)) > angSpeedTresh)
	  da(i) = angSpeedTresh * da(i) /  fabs(da(i));
    }

  // PID for the angles that are not defined in the model :
  
  if(njoints > this->seg_num)
    {
      da.Resize(njoints);
      for(int i=this->seg_num;i<njoints;i++)
	{
	  da(i) = (refPosture(i) - angles(i)) * 0.05;
	}
    }
  //angles = angles + da;
  return da;
}


Vector IKModel::inverse_step(Vector currentJoints,Vector3 target,
			     float linSpeedThresh,float angSpeedTresh)
{
  Matrix J,Jinv,Jtmp;
  Matrix Identity;
  Matrix Weight(seg_num,seg_num);
  Matrix Weight_inv(seg_num,seg_num);
  Vector3 x, dx;
  Vector da,angles=currentJoints;
  float prev_g[seg_num];
  Identity.Resize(this->seg_num,this->seg_num);

  for(int i=0;i<this->seg_num;i++) 
    {
      Identity(i,i) = 1;
      prev_g[i] = segs[i]->getJointLimitGradient(currentJoints(i));
    }

  x = direct(angles);
  dx = (target-x);
  float norm = dx.Norm();
  // thresholding dx norm ... 
  if(norm > linSpeedThresh)
    dx = dx*(linSpeedThresh/norm);
  J = this->jacobian(angles); 
  //J.Print();
  for(int i=0;i<seg_num;i++)
    {
      float g = segs[i]->getJointLimitGradient(angles(i));
      if((fabs(g)-prev_g[i]) > 0)
	{
	  Weight(i,i) = 1 + fabs(g); 
	  Weight_inv(i,i) = 1/(1+fabs(g));
	}
      else 
	{
	  Weight(i,i) = 1; 
	  Weight_inv(i,i) = 1;
	}
      prev_g[i] = fabs(g);
    }
  /* Weighted Pseudo Inverse */ 
  Jtmp = J*Weight_inv*J.Transpose(); // + damping .. lambda*Identity .. 
  Jinv = Weight_inv *J.Transpose()*Jtmp.Inverse();
  //Jinv.Print();
  da = Jinv*dx;
  if(da.Norm() > angSpeedTresh )
    da = da*(angSpeedTresh/da.Norm());
  angles = angles+da;

  /* hard threshold to avoid joint limits */
  for(int i=0;i<seg_num;i++)
    {
      if(angles(i) > segs[i]->max_joint_angle)
	angles(i) = segs[i]->max_joint_angle;
      else if(angles(i) < segs[i]->min_joint_angle)
	 angles(i) = segs[i]->min_joint_angle;
    }
  
  return angles;
}


int IKModel::inverse(Vector start,Vector3 target,Vector * next_angle)
{
  Vector3 x;
  Vector angles=start;
  x = direct(start);
  Vector3 dx;
  Vector da;
  Vector3 ctarget;
  Vector current_x;
  dx = target - x;
  float alpha;
  float norm = dx.Norm() ;
  std::cout << norm << std::endl;
  int n_iter = (int)floor(norm) + 1;

  for(int i = 1;i<= n_iter;i++)
    {
      alpha = (float) i/ (float) n_iter;
      ctarget = target*alpha + x*(1-alpha);
      angles = inverse_step(angles,ctarget,5,5,reference);
    }
  x = direct(angles);
  dx = target-x;
  norm = dx.Norm();
  std::cout << n_iter << " :: " << norm << std::endl;
  if(norm > 1) 
    return 0;
  *next_angle = angles;
  return 1;
}

void  IKModel::plotPosition()
{
  for(int j=0;j<seg_num+1;j++)
    {
      for(int k=0;k<3;k++) *f << this->joint_positions[j](k) << "  ";
      *f << std::endl;
    }
  *f << std::endl;
  *f << std::endl;
}
